import tensorflow as tf
from tensorflow import keras

from tensorflow.keras.layers import Conv2D, MaxPooling2D, Dense, Flatten, InputLayer
from tensorflow.keras.optimizers import Adam
import numpy as np
import time


# 共 22 层
class Vgg16Model():
    # your Sequential model is here.
    def __init__(self, nb_classes=1000, input_shape=(224, 224, 3)):
        self.model = keras.Sequential()
        self.model.add(InputLayer(input_shape=input_shape))

        self.model.add(Conv2D(64, (3, 3), activation='relu', padding='same', name='conv1a'))
        self.model.add(Conv2D(64, (3, 3), activation='relu', padding='same', name='conv1b'))
        self.model.add(MaxPooling2D((2, 2), strides=(2, 2), name='pool1'))

        self.model.add(Conv2D(128, (3, 3), activation='relu', padding='same', name='conv2a'))
        self.model.add(Conv2D(128, (3, 3), activation='relu', padding='same', name='conv2b'))
        self.model.add(MaxPooling2D((2, 2), strides=(2, 2), name='pool2'))

        self.model.add(Conv2D(256, (3, 3), activation='relu', padding='same', name='conv3a'))
        self.model.add(Conv2D(256, (3, 3), activation='relu', padding='same', name='conv3b'))
        self.model.add(Conv2D(256, (3, 3), activation='relu', padding='same', name='conv3c'))
        self.model.add(MaxPooling2D((2, 2), strides=(2, 2), name='pool3'))

        self.model.add(Conv2D(512, (3, 3), activation='relu', padding='same', name='conv4a'))
        self.model.add(Conv2D(512, (3, 3), activation='relu', padding='same', name='conv4b'))
        self.model.add(Conv2D(512, (3, 3), activation='relu', padding='same', name='conv4c'))
        self.model.add(MaxPooling2D((2, 2), strides=(2, 2), name='pool4'))

        self.model.add(Conv2D(512, (3, 3), activation='relu', padding='same', name='conv5a'))
        self.model.add(Conv2D(512, (3, 3), activation='relu', padding='same', name='conv5b'))
        self.model.add(Conv2D(512, (3, 3), activation='relu', padding='same', name='conv5c'))
        self.model.add(MaxPooling2D((2, 2), strides=(2, 2), name='pool5'))

        self.model.add(Flatten())
        self.model.add(Dense(4096, activation='relu', name='fc6'))
        self.model.add(Dense(4096, activation='relu', name='fc7'))

        self.model.add(Dense(nb_classes, activation='softmax', name='fc8'))

    # train is inhrent from keras.Model
    def train(self, dataset, batch_size=32, nb_epoch=50):
        ''' train function '''
        # sgd = SGD(lr=0.01, decay=1e-6,
        #           momentum=0.9, nesterov=True)  # 采用SGD+momentum的优化器进行训练，首先生成一个优化器对象
        adam = Adam(lr=1e-4)
        self.model.compile(loss='categorical_crossentropy',
                           optimizer=adam,
                           metrics=['accuracy'])  # 完成实际的模型配置工作

        self.model.fit(dataset.train_images,
                       dataset.train_labels,
                       steps_per_epoch=dataset.train_images.shape[0],
                       epochs=nb_epoch,
                       validation_data=(dataset.valid_images, dataset.valid_labels),
                       shuffle=True)

    def save_model(self, file_path='./model/vgg16model.h5'):
        self.model.save(file_path)

    def load_model(self, file_path='./model/vgg16model.h5'):
        self.model = keras.Model.load_model(file_path)

    # x is input data with shape[None, 224, 224, 3] , return a numpy darray
    def splitpredict(self, x, startlayer, endlayer=0):
        if endlayer == 0:
            endlayer = len(self.model.layers)
        if startlayer < 0 or endlayer > len(self.model.layers):
            raise Exception("Layer range wrong, please check")
        for i in range(startlayer, endlayer):
            x = self.model.layers[i](x)
        return x.numpy()

    def estimateTimeAndData(self):
        nums = 100  # 默认执行100次取平均值
        len_model = len(self.model.layers)
        time_layer = [0.0] * len_model
        datasize = [0.0] * len_model
        for i in range(100):
            x = tf.random.normal([1, 244, 244, 3], dtype=tf.float32)
            for d in range(len_model):
                start = time.process_time()
                x = self.model.layers[d](x)
                end = time.process_time()
                time_layer[d] += (end - start)
                datasize[d] += (tf.size(x).numpy())
        return np.array(time_layer) / 100, np.array(datasize) / 100


if __name__ == '__main__':
    x = tf.random.normal([1, 224, 224, 3], dtype=tf.float32)
    vgg16 = Vgg16Model(1000, (224, 224, 3))
    # vgg16.model.summary()
    # avg_time, avg_data = vgg16.estimateTimeAndData()
    # print(tf.size(x).numpy(), avg_data)
    # print(sum(avg_time))
# for i in range(len(vgg16.model.layers)):
#     print('The {} layer\'s time cost is {}, output data size is {}'.format(i+1, avg_time[i], avg_data[i]))

# print("test the splitpredict")
# index = 13
# x = vgg16.splitpredict(x, 0, index)
# print('执行到第{}层的结果:{}'.format(index, x.size))
# x = vgg16.splitpredict(x, index)
# print('继续执行后续层结果:{}'.format(x))
